package main

import (
	"fmt"
	"golang.org/x/time/rate"
	"log"
	"net/http"
)


// Limit 中间件
func Limit(next http.Handler) http.Handler {
	const (
		speed = 1  //每秒执行的次数
		capacity = 5 //桶的容量大小
	)
	var gameScene = rate.NewLimiter(speed , capacity)// 创建桶
	return http.HandlerFunc(func(writer http.ResponseWriter, request *http.Request) {
		if !gameScene.Allow() {
			http.Error(writer, "too many request", http.StatusTooManyRequests)
			return
		}
		next.ServeHTTP(writer, request)
	})
}

func main() {
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(writer http.ResponseWriter, _ *http.Request) {
		_, _ = writer.Write([]byte("ok!!!"))
	})

	addr := ":8080"
	log.Println(fmt.Sprintf("http://localhost%s", addr))
	_ = http.ListenAndServe(addr, Limit(mux))
}
